package query

import (
	"fmt"
	"strconv"
	"strings"
)

// KunlunQuery Kunlun GPU查询实现
type KunlunQuery struct{}

func init() {
	RegisterQueryer("kunlun", func() GPUQuery { return &KunlunQuery{} })
}

func (q *KunlunQuery) Query() ([]GPUInfo, error) {
	// 获取设备信息
	output, err := runCommand("xpu_smi", "-m")
	if err != nil {
		return nil, fmt.Errorf("获取设备信息失败: %v", err)
	}

	return q.parseMachineReadableOutput(output)
}

// parseMachineReadableOutput 解析机器可读格式的输出
func (q *KunlunQuery) parseMachineReadableOutput(output string) ([]GPUInfo, error) {
	lines := strings.Split(output, "\n")
	var gpus []GPUInfo

	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) < 32 {
			continue
		}

		// 解析设备索引 (字段2)
		index, err := strconv.Atoi(fields[2])
		if err != nil {
			continue
		}

		// 解析使用率 (字段19)
		utilization, err := strconv.Atoi(fields[19])
		if err != nil {
			continue
		}

		// 解析内存使用量 (字段17)
		memoryUsed, err := strconv.Atoi(fields[17])
		if err != nil {
			continue
		}

		// 解析内存总量 (字段18)
		memoryTotal, err := strconv.Atoi(fields[18])
		if err != nil {
			continue
		}

		gpus = append(gpus, GPUInfo{
			Index:       index,
			Utilization: utilization,
			MemoryUsed:  memoryUsed,
			MemoryTotal: memoryTotal,
		})
	}

	if len(gpus) == 0 {
		return nil, fmt.Errorf("未找到有效的GPU信息")
	}

	return gpus, nil
}
